
import os
import sys
import time
import md5
from feature import createFeatureParser

import snpfeature
import dbsnpfeature
import sortedsummaryfeature
import fastafeature
import wigglefeature
import coveragefeature
# import fragmentcoveragefeature
# import readsperbasecountfeature
import codonfeature
import cosmicfeature
import simplefastafeature

FILENAME_KEY = "FILENAME"
TIMESTAMP_KEY = "TIMESTAMP"
INTERVAL_KEY = "INTERVAL"
COUNT_KEY = "FEATURE_COUNT"
FEATURETYPE_KEY = "FEATURE_TYPE"
KEYS = [FILENAME_KEY, TIMESTAMP_KEY, INTERVAL_KEY, COUNT_KEY, FEATURETYPE_KEY]

INVALID_FEATURE = "INVALID_FEATURE"

class FeatureIndex:
  '''
  Represents an index for a file containing features sorted by their start position.
  Provides methods to validate an index against its source file, write an index
  to a file, read an index from a file, and retrieve all the features that overlap
  a particular location.
  '''

  def __init__(self, indexFilename=None):
    self.ixByChr = {}
    self.chrList = []
    self.indexFilename = ""
    if indexFilename != None:
      self.loadIndex(indexFilename)

  def getFeatureCount(self):
    return long(self.ixByChr[COUNT_KEY])

  def createFromFeatures(self, filename, interval=10000):
    '''Build this index from the specified file.'''
    featureParser = createFeatureParser(filename)
    n = 0
    t0 = time.time()
    currChr = None
    self.ixByChr = {}
    self.chrList = []
    self.indexFilename = ""
    self.ixByChr[FILENAME_KEY] = filename
    self.ixByChr[INTERVAL_KEY] = interval
    self.ixByChr[TIMESTAMP_KEY] = time.asctime().replace(" ", "_").replace(":", "_")
    if featureParser.getFeatureType():
      self.ixByChr[FEATURETYPE_KEY] = featureParser.getFeatureType()
    f = open(filename, "r")
    pos = 0
    for line in f:
      if line == "": break
      feature = featureParser.parse(line)
      if feature == None:
        pos += len(line)
        continue
      if currChr != feature.chr:
        currChr = feature.chr
        self.chrList.append(currChr)
        self.ixByChr[currChr] = self.ixByChr.get(currChr, [])
        self.ixByChr[currChr].append("%s\t%ld\t%ld\t%s\t%s" % (currChr, -1, pos, "###", featureParser.getLastIndexContext()))
      n += 1
      if 0 == (n % interval):
        self.ixByChr[currChr].append("%s\t%ld\t%ld\t%s\t%s" % (currChr, feature.spos, pos, md5.new(line.strip()).hexdigest(), featureParser.getLastIndexContext()))
        if 0 == (n % (10 * interval)):
          print >> sys.stderr, "# Indexed %ld lines in %.7f sec, or %.7f sec/Mfeature" % (n, time.time() - t0, 1000000.0 * (time.time() - t0) / (n + 1))
      pos += len(line)
    f.close()
    self.ixByChr[COUNT_KEY] = n
    print >> sys.stderr, "# Indexed %ld lines in %.7f sec" % (n, time.time() - t0)

  def printIndex(self, stream=sys.stdout):
    '''Print this index to the specified output stream.'''
    for key in KEYS:
      if self.ixByChr.has_key(key):
        stream.write("%s\t%s\n" % (key, self.ixByChr[key]))
    for chr in self.chrList:
      for entry in self.ixByChr.get(chr, []):
        stream.write(entry + "\n")

  def loadIndex(self, indexFilename):
    '''Load an already-generated index from a file.'''
    foundHeader = False
    self.ixByChr = {}
    self.chrList = []
    self.indexFilename = os.path.realpath(indexFilename)
    currChr = None
    f = open(self.indexFilename, "r")
    for line in f:
      cols = line.strip().split("\t")
      if len(cols) == 2 and cols[0] in KEYS:
        foundHeader = True
        self.ixByChr[cols[0]] = cols[1]
        continue
      elif len(cols) < 3:
        print >> sys.stderr, "# Warning: ignoring uninterpreted line", cols
        continue
      if not foundHeader:
        raise Exception("missing header in index file", self.indexFilename)
      (featureChr, featurePos, pos) = (cols[0], cols[1], long(cols[2]))
      if currChr != featureChr:
        currChr = featureChr
        self.chrList.append(currChr)
        self.ixByChr[currChr] = self.ixByChr.get(currChr, [])
      self.ixByChr[currChr].append(line.strip())
    f.close()

  def validationError(self):
    '''
    Attempt to validate this index against its source file.  Return a string
    describing the error, if one occurs, or None if validation is successful.
    '''
    if not self.ixByChr.has_key(FILENAME_KEY):
      return "Missing source filename"
    indexDir = os.path.dirname(self.indexFilename)
    featureFilename = os.path.join(indexDir, self.ixByChr[FILENAME_KEY])
    f = open(featureFilename, "r")
    for chr in self.chrList:
      for entry in self.ixByChr[chr]:
        entryCols = entry.strip().split("\t")
        minPos = long(entryCols[1])
        pos = long(entryCols[2])
        cksum = entryCols[3]
        f.seek(pos)
        line = f.readline()
        if minPos >= 0 and cksum != md5.new(line.strip()).hexdigest():
          f.close()
          return "Bad index entry:\n%s\n%s" % (entry, line)
    f.close()
    return None

  def getFeatures(self, query=(None, None, None)):
    '''
    Return an iterator over features overlapping the specified
    query interval, given in the form (chr, startPos, endPos).
    Any of chr, startPos, and endPos may be given as None to
    drop the corresponding restriction; e.g., (None, None, None)
    returns all features.
    '''
    (qchr, qmin, qmax) = query
    if qchr == None:
      # if None, iterate over all chromosomes in index
      for chr in self.chrList:
        for feature in self.getFeatures((chr, qmin, qmax)):
          yield feature
      return
    if not self.ixByChr.has_key(FILENAME_KEY):
      yield INVALID_FEATURE
      return
    indexDir = os.path.dirname(self.indexFilename)
    featureFilename = os.path.join(indexDir, self.ixByChr[FILENAME_KEY])
    featureParser = createFeatureParser(featureFilename)
    maxFeatureSize = featureParser.getMaxFeatureSize()
    # print >> sys.stderr, "# Performing query", query
    ix = self.ixByChr.get(qchr, [])
    f = open(featureFilename, "r")
    lastPos = -1
    for entry in ix:
      entryCols = entry.strip().split("\t")
      entryMinPos = long(entryCols[1])
      entryFilePos = long(entryCols[2])
      entryCksum = entryCols[3]
      entryContext = ""
      if len(entryCols) > 4:
        entryContext = entryCols[4]
      if entryMinPos < 0 or (maxFeatureSize != None and qmin != None and entryMinPos + maxFeatureSize <= qmin):
        lastPos = entryFilePos
        lastContext = entryContext
      else:
        break
    if lastPos >= 0:
      f.seek(lastPos)
      featureParser.setIndexContext(lastContext)
    for line in f:
      if line == "": break
      feature = featureParser.parse(line)
      if feature == None: continue
      (chr, minPos, maxPos) = (feature.chr, feature.spos, feature.epos)
      if chr != qchr or (qmax != None and minPos > qmax): break
      if (qmin != None and maxPos < qmin): continue
      yield feature
    f.close()

def validateIndex(filename):
  index = FeatureIndex(filename)
  err = index.validationError()
  if err != None:
    raise Exception(err)
  
def getFeatures(filename, query=(None, None, None)):
  '''Read features from an indexed file, falling back to unindexed file if needed.'''
  if not filename.endswith("INDEX"):
    print >> sys.stderr, filename, "does not end with INDEX; attempting (slow) read of unindexed features..."
    features = _getUnindexedFeatures(filename, query)
  else:
    try:
      index = FeatureIndex(filename)
    except IOError, e:
      yield INVALID_FEATURE
      return
    features = index.getFeatures(query)
  for feature in features:
    yield feature
  return

def _getUnindexedFeatures(filename, query=(None, None, None)):
  '''
  Fallback method for reading features from original, unindexed file.
  Query given as (chr, startPos, endPos); if any of these are None,
  the corresponding restriction is droppped.
  '''
  featureParser = createFeatureParser(filename)
  (qchr, qmin, qmax) = query
  f = open(filename, "r")
  foundBlock = False
  for line in f:
    feature = featureParser.parse(line)
    if feature == None: continue
    (chr, minPos, maxPos) = (feature.chr, feature.spos, feature.epos)
    if (qchr != None and chr != qchr) or (qmax != None and minPos > qmax):
      if foundBlock: break
      continue
    if (qmin != None and maxPos < qmin): continue
    foundBlock = True
    yield feature
  f.close()

if __name__ == "__main__":
  if sys.argv[1] == "-I":
    filename = sys.argv[2]
    index = FeatureIndex()
    interval = 10000
    if len(sys.argv) > 3:
      interval = int(sys.argv[3])
      print >> sys.stderr, "# Indexing every %d lines..." % interval
    index.createFromFeatures(filename, interval)
    index.printIndex()
  else:
    indexFilename = sys.argv[1]
    if len(sys.argv) < 4:
      validateIndex(indexFilename)
    else:
      minPos = maxPos = long(sys.argv[3])
      if len(sys.argv) > 4: maxPos = long(sys.argv[4])
      query = (sys.argv[2], minPos, maxPos)
      features = getFeatures(indexFilename, query)
      for f in features:
        print (f.chr, f.spos, f.epos, f.content.strip())

